# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/12 13:38
# @Function:
import glob

from osgeo import gdal
from rasterio.errors import RasterioIOError
# from rasterio.enums import Resampling
from rasterio.warp import reproject, Resampling
from rasterio.windows import Window

import os
import rasterio
import numpy as np
import warnings


# 忽略特定的 RasterioDeprecationWarning
warnings.filterwarnings("ignore", category=FutureWarning, message="statistics() will be removed in 2.0.0.*")

def replace_path_segment(original_path: str, old_segment: str, new_segment: str) -> str:
    """
    替换路径中的某个段落。

    参数:
        original_path (str): 原始路径。
        old_segment (str): 要替换的旧段落。
        new_segment (str): 新的段落。

    返回:
        str: 替换后的路径。
    """
    # 使用 os.path.normpath 规范化路径，确保路径分隔符统一
    normalized_path = os.path.normpath(original_path)

    # 使用 os.path.split 将路径分割为目录部分和最后一级文件夹/文件名
    head, tail = os.path.split(normalized_path)

    # 替换最后一级文件夹/文件名中的指定段落
    new_tail = tail.replace(old_segment, new_segment)

    # 使用 os.path.join 重新组合路径
    new_path = os.path.join(head, new_tail)
    os.makedirs(new_path, exist_ok=True)  # 忽略路径已存在的情况
    print(f"创建New Path: {new_path}")
    # if not os.path.exists(new_path):
    #     os.makedirs(new_path)
    #     print(f"创建New Path: {new_path}")

    return new_path

def cut_tif_to_chunks(input_tif_path, num_chunks=40, output_folder=None):
    """
    将输入的 TIF 文件切成num_chunks个tif文件， 切割后的TIF 文件保存了金字塔、压缩、统计信息。

    参数:
        input_path (str): 输入的 TIF 文件路径。
        num_chunks (int): 想要将TIF 文件切成多少个块数，默认为 40。
        output_path (str): 输出的 TIF 文件路径。默认是在输入文件夹中创建新文件夹，命名规范: TIF文件名_chunks


    返回:
        None
    """
    # 读取输入的 TIFF 文件
    with rasterio.open(input_tif_path) as src:
        height = src.height
        width = src.width
        num_bands = src.count

        # 计算每个切割块的大小，使其近似为矩形
        num_cols = int(np.ceil(np.sqrt(num_chunks)))
        num_rows = int(np.ceil(num_chunks / num_cols))
        chunk_height = height // num_rows
        chunk_width = width // num_cols

        # 如果指定了输出文件夹，则使用指定的文件夹，否则在输入文件夹中创建一个新文件夹
        if output_folder is None:
            input_folder = os.path.dirname(input_tif_path)
            input_filename = os.path.basename(input_tif_path)
            output_folder = os.path.join(input_folder, f"{os.path.splitext(input_filename)[0]}_chunks")

        # 确保输出文件夹存在
        os.makedirs(output_folder, exist_ok=True)

        # 切割并保存文件
        chunk_index = 1
        for row in range(num_rows):
            for col in range(num_cols):
                if chunk_index > num_chunks:
                    break

                # 计算当前块的窗口
                row_off = row * chunk_height
                col_off = col * chunk_width
                window = Window(col_off, row_off, chunk_width, chunk_height)

                # 读取当前块的数据
                chunk_data = src.read(window=window)

                # 计算统计信息
                mean_value = np.mean(chunk_data, axis=(1, 2))
                std_value = np.std(chunk_data, axis=(1, 2))

                # 构建输出文件路径
                output_filename = f"{os.path.splitext(input_filename)[0]}_chunk_{chunk_index}.tif"
                output_path = os.path.join(output_folder, output_filename)

                # 写入输出文件，使用 LZW 压缩
                with rasterio.open(
                    output_path,
                    'w',
                    driver='GTiff',
                    height=chunk_height,
                    width=chunk_width,
                    count=num_bands,
                    dtype=src.dtypes[0],
                    crs=src.crs,
                    transform=src.window_transform(window),
                    compress='lzw'  # 使用 LZW 压缩
                ) as dst:
                    dst.write(chunk_data)

                    # # 添加统计信息到元数据
                    # metadata = {
                    #     'mean': mean_value.tolist(),
                    #     'std': std_value.tolist()
                    # }
                    # dst.update_tags(**metadata)

                    # 添加金字塔信息
                    overviews_levels = [2, 4, 8, 16]  # 金字塔级别
                    dst.build_overviews(overviews_levels, Resampling.average)

                    # 显式生成统计信息（关键修复）
                    dst.write_colormap(1, {})  # 清空伪彩色表
                    dst.update_tags(ns='rio_stats', STATISTICS_VALID_PERCENT=100)  # 标记有效统计
                    # 强制重新计算统计信息（排除无效值）
                    # stats = dst.stats(1)
                    stats = dst.statistics(1)  # 或者使用正确的属性名
                    print(f"实际数据范围: {stats.min} ~ {stats.max}")  # 验证输出范围

                print(f"Saved {output_path} with mean={mean_value} and std={std_value}")

                chunk_index += 1

# 示例调用
# input_tif_path = r"F:\test_files\插值验证程序的输入数据\双线性插值\et_2023.tif"
# # output_folder = "path/to/your/output/folder"  # 可选参数，如果不指定则在输入文件夹中创建一个新文件夹
# cut_tif_to_chunks(input_tif_path, num_chunks=40)

# 添加了金字塔和压缩
def interpolate_to_target_resolution_01(input_path: str, target_resolution: float = 0.00027777778, output_path: str = None) -> None:
    """
    01:添加了金字塔、压缩、统计信息
    将输入的 TIF 文件插值到目标分辨率（默认 0.01°），并保存结果。

    参数:
        input_path (str): 输入的 TIF 文件路径。
        output_path (str): 输出的 TIF 文件路径。
        target_resolution (float): 目标分辨率（单位：度），默认为 0.01°。

    返回:
        None
    """
    # 自动生成输出路径
    if output_path is None:
        input_tif_folder = os.path.dirname(input_path)
        output_folder = replace_path_segment(input_tif_folder, "chunks", "interpolation")

        ext = os.path.splitext(input_path)[1]
        meters = int(target_resolution * 111000)  # 1度≈111km

        input_filename = os.path.basename(input_path)
        input_filename_without_ext = os.path.splitext(input_filename)[0]
        output_path = f"{output_folder}\\{input_filename_without_ext}_{meters}m{ext}"


    try:
        # 打开输入文件
        with rasterio.open(input_path) as src:
            # 读取数据
            data = src.read(1)
            transform = src.transform
            crs = src.crs
            nodata = src.nodata  # 获取无效值

            # 判断 CRS 是否缺失
            if crs is None:
                print("警告：输入文件未包含有效的 CRS，将使用默认 CRS (EPSG:4326/WGS84)。")
                crs = "EPSG:4326"  # 设置默认 CRS 为 WGS84

            # 如果 nodata 值不是 -9999，将其修改为 -9999
            if nodata is not None and nodata != -9999:
                data[data == nodata] = -9999  # 替换 nodata 值为 -9999
                nodata = -9999.0  # 更新 nodata 值

            # 将无效值替换为 nodata
            data = data.astype('float32')  # 转换为浮点型以支持 NaN
            nodata = float(nodata) if nodata is not None else -9999.0  # 统一为浮点数
            # 使用浮点数阈值判断（1e-6为误差容忍度）
            mask = np.isclose(data, nodata, atol=1e-6, rtol=0)
            data[mask] = np.nan

            # 计算目标尺寸
            target_width = int((src.bounds.right - src.bounds.left) / target_resolution)
            target_height = int((src.bounds.top - src.bounds.bottom) / target_resolution)

            # 创建目标数组
            target_data = np.full((target_height, target_width), nodata, dtype=data.dtype)

            # 计算目标变换矩阵
            target_transform = rasterio.Affine(
                target_resolution, 0, src.bounds.left,
                0, -target_resolution, src.bounds.top
            )

            # 使用双线性插值进行重投影
            reproject(
                source=data,
                destination=target_data,
                src_transform=transform,
                src_crs=crs,
                dst_transform=target_transform,
                dst_crs=crs,
                resampling=Resampling.bilinear,  # 使用双线性插值
                src_nodata=nodata,  # 指定源数据的无效值
                dst_nodata=nodata  # 指定目标数据的无效值
            )

            # 将 NaN 替换回原始无效值（如果存在）
            target_data[np.isnan(target_data)] = nodata

            # 保存插值结果
            with rasterio.open(
                    output_path,
                    'w',
                    driver='GTiff',
                    height=target_height,
                    width=target_width,
                    count=1,
                    dtype=target_data.dtype,
                    crs=crs,
                    transform=target_transform,
                    nodata=nodata,  # 设置输出文件的无效值
                    compress='lzw'  # 使用 LZW 压缩
            ) as dst:
                dst.write(target_data, 1)

                # 添加金字塔信息
                overviews_levels = [2, 4, 8, 16]  # 金字塔级别
                dst.build_overviews(overviews_levels, Resampling.average)

                # 显式生成统计信息（关键修复）
                dst.write_colormap(1, {})  # 清空伪彩色表
                dst.update_tags(ns='rio_stats', STATISTICS_VALID_PERCENT=100)  # 标记有效统计
                # 强制重新计算统计信息（排除无效值）
                # stats = dst.stats(1)
                stats = dst.statistics(1)
                print(f"实际数据范围: {stats.min} ~ {stats.max}")  # 验证输出范围

    except RasterioIOError as e:
        print(f"文件 {input_path} 打开失败: {e}")
    except Exception as e:
        print(f"未知错误: {e}")

    print(f"Interpolated data saved to {output_path}")

# 示例调用
# input_path = r'F:\test_files\插值验证程序的输入数据\双线性插值\2023年降水量空间分布栅格图.tif'
# output_path = r'F:\test_files\插值验证程序的输入数据\双线性插值\2023年降水量空间分布栅格图_30m.tif'
# # 插值到 0.00027777778° 分辨率
# interpolate_to_target_resolution(input_path, output_path, target_resolution=0.0027777778)


# input_path = r'F:\test_files\插值验证程序的输入数据\双线性插值\et_2023_chunks\chunk_2.tif'
# # output_path = r'F:\test_files\插值验证程序的输入数据\双线性插值\2023年降水量空间分布栅格图_30m.tif'
# # 插值到 0.00027777778° 分辨率
# interpolate_to_target_resolution_01(input_path, target_resolution=0.00027777778)


def merge_tif_files_for_one(input_path, merge_name: str = None):
    """
    合并文件夹下的所有tif文件为一个文件。
    输出文件保存在输入的文件夹下创建的merged文件夹中
    参数:
    input_path -- 输入文件夹的路径，包含所有tif文件。
    merge_name -- 输出合并文件的名称。
    """
    if merge_name is None:
        new_path = replace_path_segment(input_path, "interpolation", "")
        input_filename = os.path.basename(new_path)
        merge_name = os.path.splitext(input_filename)[0]

    # 获取文件夹下的所有tif文件
    input_files = glob.glob(os.path.join(input_path, '*.tif'))
    print(f"查找到{input_path}文件夹中{len(input_files)} 个tif文件。")

    if not input_files:
        print("没有找到任何tif文件。")
        return

    # 存储有效的tif文件
    inputrasfile = []

    for fn in input_files:
        infn = gdal.Open(fn, gdal.GA_ReadOnly)
        if infn is None:
            print(f"文件 {fn} 无法打开，跳过此文件。")
            continue
        input_proj1 = infn.GetProjection()
        if input_proj1 is None:
            print(f"文件 {fn} 没有有效的投影信息，跳过此文件。")
            continue
        inputrasfile.append(infn)

    # 如果没有有效的输入文件，退出函数
    if not inputrasfile:
        print("没有有效的输入文件，无法进行合并。")
        return

    # 确保输出文件夹存在
    output_file = os.path.join(input_path, 'merged', merge_name + 'merged.tif')
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    # 使用gdal.Warp进行合并
    # 设置输出参数，可以指定投影信息、目标分辨率、合并方法等
    options = gdal.WarpOptions(
        format='GTiff',
        creationOptions=[
            'COMPRESS=LZW',  # 使用 LZW 压缩
            'TILED=YES',  # 启用瓦片
            'BIGTIFF=IF_SAFER'  # 如果需要，启用大文件支持
        ],
        dstNodata=-9999,  # 设置无效值
    )
    gdal.Warp(output_file, inputrasfile, options=options)

    # 重新打开合并后的文件以添加金字塔信息和统计信息
    with rasterio.open(output_file, 'r+') as dst:
        # 添加金字塔信息
        overviews_levels = [2, 4, 8, 16]  # 金字塔级别
        dst.build_overviews(overviews_levels, Resampling.average)

        # 显式生成统计信息（关键修复）
        dst.write_colormap(1, {})  # 清空伪彩色表
        dst.update_tags(ns='rio_stats', STATISTICS_VALID_PERCENT=100)  # 标记有效统计
        # 强制重新计算统计信息（排除无效值）
        stats = dst.statistics(1)
        print(f"实际数据范围: {stats.min} ~ {stats.max}")  # 验证输出范围

    print(f"合并完成，文件保存为: {output_file}")